#include <stdlib.h>
#include <stdio.h>

#include "relic.h"
#include "sqft_conf.h"
#include "relic_test.h"


int perf_tests(void) {
	int code = STS_ERR;
	bn_t a, b, c, d, e, u, m;
	
	bn_null(a);
	bn_null(b);
	bn_null(c);
	bn_null(d);
	bn_null(e);
	bn_null(u);
	bn_null(m);

	TRY {
		bn_new(a);
		bn_new(b);
		bn_new(c);
		bn_new(d);
		bn_new(e);
		bn_new(u);
		bn_new(m);
		
		bn_gen_prime(m, BN_BITS/2);
		bn_mod_pre(u, m);
		/*
		TEST_BEGIN("test mod monty") {
			bn_rand(a, BN_POS, BN_BITS/2);
			bn_rand(c, BN_POS, BN_BITS/2);
			bn_mod_monty_conv(b, a, m);
			
			
			bn_sqr(a, a);
			bn_mod(a, a, m);
			bn_add(a, a, a);
			if (bn_cmp(a, m) != CMP_LT) {
				bn_sub(a, a, m);
			}
			
			bn_sqr(b, b);
			bn_mod(b, b, m, u);
			bn_add(b, b, b);
			if (bn_cmp(b, m) != CMP_LT) {
				bn_sub(b, b, m);
			}
			
			bn_mod_monty_back(b, b, m);
			
			TEST_ASSERT(bn_cmp(a, b) == CMP_EQ, end);
		}
		TEST_END;
		*/
		TEST_BEGIN("test zero power is correct") {
			bn_rand(a, BN_POS, BN_BITS);
			bn_zero(b);
			bn_set_dig(d, 1);
			bn_mxp(c, a, b, m);
			TEST_ASSERT(bn_cmp(c, d) == CMP_EQ, end);
		}
		TEST_END;
		

	}
	CATCH_ANY {
		ERROR(end);
	}
	code = STS_OK;
	end:
	bn_free(a);
	bn_free(b);
	bn_free(c);
	bn_null(d);
	bn_null(e);
	bn_free(u);
	bn_free(m);
	return code;
}


int util(void) {
	int code = STS_ERR;
	bn_t a, b;
	
	bn_null(a);
	bn_null(b);

	TRY {
		bn_new(a);
		bn_new(b);
		
		TEST_BEGIN("setting signed digit is correct") {
			bn_set_dis(a, -1);
			TEST_ASSERT(bn_is_minus_one(a), end);
			bn_set_dis(a, 1);
			TEST_ASSERT(bn_cmp_dig(a, 1) == CMP_EQ, end);
		}
		TEST_END;
		
		TEST_BEGIN("square root is consistent") {
			bn_rand(a, BN_POS, BN_BITS/2);
			bn_sqr(b, a);
			bn_sqt(b, b);
			TEST_ASSERT(bn_cmp(a, b) == CMP_EQ, end);
		}
		TEST_END;
		
		TEST_BEGIN("perfect square test is correct") {
			bn_rand(a, BN_POS, BN_BITS/2);
			bn_sqr(a, a);
			bn_sub_dig(b, a, 2);
			TEST_ASSERT(bn_is_perfect_sqr(a), end);
			TEST_ASSERT(!bn_is_perfect_sqr(b), end);
		}
		TEST_END;
	}
	CATCH_ANY {
		ERROR(end);
	}
	code = STS_OK;
	end:
	bn_free(a);
	bn_free(b);
	return code;
}

int pol(void) {
	int ep, code = STS_ERR;
	bn_t a[2], b[2], c[2], d[2], e[2], f[2], m[2];
	bn_t t, p;
	
	bn_pol_null(a);
	bn_pol_null(b);
	bn_pol_null(c);
	bn_pol_null(d);
	bn_pol_null(e);
	bn_pol_null(f);
	bn_pol_null(m);
	bn_null(t);
	bn_null(p);

	TRY {
		bn_pol_new(a);
		bn_pol_new(b);
		bn_pol_new(c);
		bn_pol_new(d);
		bn_pol_new(e);
		bn_pol_new(f);
		bn_pol_new(m);
		bn_new(t);
		bn_new(p);
		
		
		
		
		/*---------------util-----------------*/
		
		bn_gen_prime(p, BN_BITS);
		bn_pol_zero(m);
		bn_rand(m[0], BN_POS, bn_bits(p));
		bn_mod(m[0], m[0], p);
		
		TEST_BEGIN("polynomial comparison is consistent") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_rand_in_ring(b, p);
			if (bn_pol_cmp(a, b) != CMP_EQ) {
				if (bn_pol_cmp(a, b) == CMP_GT) {
					TEST_ASSERT(bn_pol_cmp(b, a) == CMP_LT, end);
				} else {
					TEST_ASSERT(bn_pol_cmp(b, a) == CMP_GT, end);
				}
			}
			TEST_ASSERT(bn_pol_cmp(a, a) == CMP_EQ, end);
		}
		TEST_END;
		
		TEST_BEGIN("assignment to polynom zero and zero test are consistent") {
			bn_pol_zero(a);
			TEST_ASSERT(bn_pol_is_zero(a), end);
		} TEST_END;
		
		
		TEST_BEGIN("polynom copy is consistent") {
			do {
				bn_pol_rand_in_ring(a, p);
				bn_pol_rand_in_ring(c, p);
			} while (bn_pol_cmp(a, c) == CMP_EQ);
			bn_pol_copy(c, a);
			TEST_ASSERT(bn_pol_cmp(c, a) == CMP_EQ, end);
		}
		TEST_END;
		
		/*---------------addition-----------------*/
		
		bn_gen_prime(p, BN_BITS);
		bn_rand(m[0], BN_POS, bn_bits(p));
		bn_mod(m[0], m[0], p);
		
		TEST_BEGIN("addition in polynomial ring is commutative") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_rand_in_ring(b, p);
			bn_pol_add(d, a, b, p);
			bn_pol_add(e, b, a, p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;

		TEST_BEGIN("addition in polynomial ring is associative") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_rand_in_ring(b, p);
			bn_pol_rand_in_ring(c, p);
			bn_pol_add(d, a, b, p);
			bn_pol_add(d, d, c, p);
			bn_pol_add(e, b, c, p);
			bn_pol_add(e, a, e, p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;

		TEST_BEGIN("addition in polynomial ring has identity") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_zero(d);
			bn_pol_add(e, a, d, p);
			TEST_ASSERT(bn_pol_cmp(e, a) == CMP_EQ, end);
			bn_pol_add(e, d, a, p);
			TEST_ASSERT(bn_pol_cmp(e, a) == CMP_EQ, end);
		} TEST_END;

		TEST_BEGIN("addition has inverse") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_neg(d, a, p);
			bn_pol_add(e, a, d, p);
			TEST_ASSERT(bn_pol_is_zero(e), end);
		} TEST_END;
		
		
		/*---------------multiplication-----------------*/
		
		
		bn_gen_prime(p, BN_BITS/2);
		bn_rand(m[0], BN_POS, bn_bits(p));
		bn_mod(m[0], m[0], p);
			
		TEST_BEGIN("multiplication in polynomial ring is commutative") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_rand_in_ring(b, p);
			bn_pol_mul(d, a, b, m, p);
			bn_pol_mul(e, b, a, m, p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;

		bn_gen_prime(p, BN_BITS/3);
		bn_rand(m[0], BN_POS, bn_bits(p));
		bn_mod(m[0], m[0], p);

		TEST_BEGIN("multiplication in polynomial ring is associative") {			
			bn_pol_rand_in_ring(a, p);
			bn_pol_rand_in_ring(b, p);
			bn_pol_rand_in_ring(c, p);
			bn_pol_mul(d, a, b, m, p);
			bn_pol_mul(d, d, c, m, p);
			bn_pol_mul(e, b, c, m, p);
			bn_pol_mul(e, a, e, m, p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;

		bn_gen_prime(p, BN_BITS/2);
		bn_rand(m[0], BN_POS, bn_bits(p));
		bn_mod(m[0], m[0], p);

		TEST_BEGIN("multiplication in polynomial ring is distributive") {			
			bn_pol_rand_in_ring(a, p);
			bn_pol_rand_in_ring(b, p);
			bn_pol_rand_in_ring(c, p);
			bn_pol_add(d, a, b, p);
			bn_pol_mul(d, c, d, m, p);
			bn_pol_mul(e, c, a, m, p);
			bn_pol_mul(f, c, b, m, p);
			bn_pol_add(e, e, f, p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;

		bn_gen_prime(p, BN_BITS);
		bn_rand(m[0], BN_POS, bn_bits(p));
		bn_mod(m[0], m[0], p);
		
		TEST_BEGIN("multiplication in polynomial ring has identity") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_set_id(d);
			bn_pol_mul(e, a, d, m, p);
			TEST_ASSERT(bn_pol_cmp(e, a) == CMP_EQ, end);
		} TEST_END;
		
		TEST_BEGIN("multiplication in polynomial ring has zero property") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_zero(d);
			bn_pol_mul(e, a, d, m, p);
			TEST_ASSERT(bn_pol_is_zero(e), end);
		} TEST_END;
		
		
		/*---------------squaring-----------------*/
		
		bn_gen_prime(p, BN_BITS/2);
		bn_set_dig(m[0], 2);
		
		TEST_BEGIN("squaring in polynomial ring is correct") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_mul(b, a, a, m, p);
			bn_pol_sqr_basic(c, a, m, p);
			TEST_ASSERT(bn_pol_cmp(b, c) == CMP_EQ, end);
			bn_pol_sqr_karat(c, a, m, p);
			TEST_ASSERT(bn_pol_cmp(b, c) == CMP_EQ, end);
		} TEST_END;
		
		
		/*---------------mod exp-----------------*/
		
		
		TEST_BEGIN("negation and mod exp in polynomial ring is correct") {
			bn_set_dig(c[0], 3);
			bn_pol_rand_in_ring(a, p);
			bn_pol_neg(b, a, p);
			bn_pol_mxp_basic(a, a, c[0], m, p);
			bn_pol_mxp_basic(e, b, c[0], m, p);
			bn_pol_neg(d, a, p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;
		
		TEST_BEGIN("exp zero is correct") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_set_id(b);
			bn_zero(t);
			bn_pol_mxp_basic(c, a, t, m, p);
			//TEST_ASSERT(bn_pol_cmp(c, b) == CMP_EQ, end);
			//bn_pol_mxp_slide(c, a, t, m, p);
			//TEST_ASSERT(bn_pol_cmp(c, b) == CMP_EQ, end);
		}
		TEST_END;
		
		bn_gen_prime(p, BN_BITS);
		bn_set_dig(c[0], 2);
		bn_mod_pre(f[0], p);

/*
		TEST_BEGIN("multiplication switch is correct") {
			bn_gen_prime(p, BN_BITS);
			bn_rand(a[0], BN_POS, bn_bits(p)/2);
			bn_set_dis(b[0], -1);
			bn_sub_dig(t, p, 1);
			bn_mul(c[0], a[0], t);
			bn_mod(c[0], c[0], p);
			if (!bn_pol_mul_swt(d[0], a[0], b[0], p)) bn_mod(d[0], d[0], p);
			TEST_ASSERT(bn_cmp(c[0], d[0]) == CMP_EQ, end);
			
			bn_set_dig(b[0], 2);
			bn_mul(c[0], a[0], b[0]);
			bn_mod(c[0], c[0], p);
			if (!bn_pol_mul_swt(d[0], a[0], b[0], p)) bn_mod(d[0], d[0], p);
			TEST_ASSERT(bn_cmp(c[0], d[0]) == CMP_EQ, end);
			
			bn_rand(b[0], BN_POS, bn_bits(p)/2);
			bn_mul(c[0], a[0], b[0]);
			bn_mod(c[0], c[0], p);
			if (!bn_pol_mul_swt(d[0], a[0], b[0], p)) bn_mod(d[0], d[0], p);
			TEST_ASSERT(bn_cmp(c[0], d[0]) == CMP_EQ, end);	
		} TEST_END;*/
		
		TEST_BEGIN("montgomery reduction in modmul is correct") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_rand_in_ring(b, p);
			bn_pol_mul_karat(d, a, b, c, p);
			bn_mod_monty_conv(a[0], a[0], p);
			bn_mod_monty_conv(a[1], a[1], p);
			bn_mod_monty_conv(b[0], b[0], p);
			bn_mod_monty_conv(b[1], b[1], p);
			bn_pol_mul_karat_u(e, a, b, c, p, f[0]);
			bn_mod_monty_back(e[0], e[0], p);
			bn_mod_monty_back(e[1], e[1], p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;
		
		TEST_BEGIN("montgomery reduction in modsqr is correct") {
			bn_pol_rand_in_ring(a, p);
			bn_pol_sqr_basic(d, a, c, p);
			bn_mod_monty_conv(b[0], a[0], p);
			bn_mod_monty_conv(b[1], a[1], p);
			bn_pol_sqr_u(e, b, c, p, f[0]);
			bn_mod_monty_back(e[0], e[0], p);
			bn_mod_monty_back(e[1], e[1], p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;
		
		
		TEST_BEGIN("mod exp in polynomial ring is correct") {
			bn_set_dig(c[0], 100);
			bn_pol_rand_in_ring(a, p);
			bn_pol_mxp_basic(d, a, c[0], m, p);
			bn_pol_mxp_slide(e, a, c[0], m, p);
			TEST_ASSERT(bn_pol_cmp(d, e) == CMP_EQ, end);
		} TEST_END;
		
		
		TEST_BEGIN("cat's jump is consistent") {
			bn_gen_prime(p, 32);
			bn_is_prime_mr2(p, c[0], f);
			do {
				bn_pol_rand_in_ring(a, p);
				bn_pol_norm(m[0], a, c, p);
				bn_smb_jac(t, m[0], p);
			} while(bn_is_minus_one(t));
			bn_sqr(t, p);
			bn_sub_dig(t, t, 1);
			bn_rsh(t, t, 3);
			bn_pol_mxp_slide(b, a, t, c, p);
			
			bn_sub_dig(t, p, 1);
			ep = t->dp[0] & 0b111;
			bn_rsh(t, t, 3);
			bn_pol_mxp_slide(d, a, t, c, p);
			bn_pol_norm(m[0], d, c, p);
			m[0]->sign = BN_POS;
			
			if (ep != 0) {
				bn_set_dig(t, (ep*(ep+2))/8);
				bn_pol_mxp_slide(e, a, t, c, p);
				
				bn_set_dig(t, ep);
				bn_pol_mxp_slide(d, d, t, c, p);
				bn_pol_mul_karat(d, d, e, c, p);
				
							
				bn_mul(d[0], d[0], m[0]);
				bn_mul(d[1], d[1], m[0]);
				
				bn_mod(d[0], d[0], p);
				bn_mod(d[1], d[1], p);
			}
			else {
				bn_zero(d[1]);
				bn_copy(d[0], m[0]);
			}
			
			TEST_ASSERT(bn_pol_cmp(b, d) == CMP_EQ, end);
		} TEST_END;

	}
	CATCH_ANY {
		ERROR(end);
	}
	code = STS_OK;
	end:
	bn_pol_free(a);
	bn_pol_free(b);
	bn_pol_free(c);
	bn_pol_free(d);
	bn_pol_free(e);
	bn_pol_free(f);
	bn_pol_free(m);
	bn_free(t);
	bn_free(p);
	return code;
}


int prime(void) {
	int bits, code = STS_ERR;
	bn_t p;
	
	bn_null(p);

	TRY {
		bn_new(p);
		TEST_BEGIN("SQFT test for compositeness is correct") {
			//generates an odd not prime number
			do {
				bn_rand(p, BN_POS, BN_BITS);
			} while (bn_is_even(p) || bn_is_prime(p));
			TEST_ASSERT(!bn_is_prime_sqft(p), end);
		}
		TEST_END;
		
		//Bug in jac rounds for 7, 8 and 9 bits.
		for (bits=8; bits<=BN_BITS; bits *= 2) {
			printf("%d BITS-----\n", bits);
			TEST_BEGIN("SQFT test for all bits is correct") {
				bn_gen_prime(p, bits);
				TEST_ASSERT(bn_is_prime_sqft(p), end);
			}
			TEST_END;
		}
	}
	CATCH_ANY {
		ERROR(end);
	}
	code = STS_OK;
	end:
	bn_free(p);
	return code;
}


int main(int argc, char *argv[]) {
	bn_t a;
	if (core_init() != STS_OK) {
		core_clean();
		return 1;
	}
	
	bn_null(a);
	bn_new(a);
	
	util_banner("Tests for SQFT components", 0);
	util_banner("Utilities:", 1);
	
	if (perf_tests() != STS_OK) {
		core_clean();
		return 1;
	}

	if (util() != STS_OK) {
		core_clean();
		return 1;
	}
	
	util_banner("Polynomial operations:", 1);
	
	if (pol() != STS_OK) {
		core_clean();
		return 1;
	}

	util_banner("SQFT prime test:", 1);
	
	if (prime() != STS_OK) {
		core_clean();
		return 1;
	}

	util_banner("All tests have passed.\n", 0);
	
	bn_free(a);
	core_clean();
	return 0;
}
